package com.yeskery.nut.extend.mcp.server;

import com.yeskery.nut.bean.ApplicationContext;
import com.yeskery.nut.core.*;
import com.yeskery.nut.extend.mcp.core.*;
import com.yeskery.nut.extend.mcp.server.method.*;
import com.yeskery.nut.extend.responsive.JsonResponsive;
import com.yeskery.nut.extend.responsive.ResponsiveConvert;
import com.yeskery.nut.extend.responsive.TypeReference;
import com.yeskery.nut.util.JsonNullRemover;
import com.yeskery.nut.util.StringUtils;

import java.util.HashMap;
import java.util.Map;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * 基础的Http MCP消息请求处理器
 * @author sprout
 * 2025-04-17 10:50
 */
public abstract class BaseHttpMcpMessageRequestHandler implements RequestHandler, ResponsiveConvertExpanded {

    /** 日志对象 */
    private static final Logger logger = Logger.getLogger(BaseHttpMcpMessageRequestHandler.class.getName());

    /** 应用上下文 */
    private final ApplicationContext applicationContext;

    /** 环境对象 */
    private final Environment environment;

    /** 响应式转换器 */
    private final ResponsiveConvert responsiveConvert;

    /** MCP会话管理器 */
    private final McpSessionManager mcpSessionManager;

    /** Mcp服务能力 */
    private final McpServerCapability mcpServerCapability;

    /** 方法处理器映射 */
    private final Map<String, MethodHandler> methodHandlerMap = new HashMap<>();

    /**
     * 构建的Http MCP消息请求处理器
     * @param applicationContext 应用上下文
     * @param mcpSessionManager MCP会话管理器
     * @param mcpServerCapability MCP服务能力
     */
    public BaseHttpMcpMessageRequestHandler(ApplicationContext applicationContext, McpSessionManager mcpSessionManager, McpServerCapability mcpServerCapability) {
        this.applicationContext = applicationContext;
        this.mcpSessionManager = mcpSessionManager;
        this.mcpServerCapability = mcpServerCapability;
        this.environment = applicationContext.getBean(Environment.class);
        this.responsiveConvert = applicationContext.getBean(JsonResponsive.class).getResponsiveConvert();
        initMethodHandler();
    }

    @Override
    public void handle(Request request, Response response, Execution execution) {
        boolean supported = supportServerSession();
        String sessionId = null;
        McpServerSession serverSession = null;
        if (supported) {
            sessionId = request.getParameter(Constants.SESSION_ID);
            if (StringUtils.isBlank(sessionId)) {
                response.setCode(ResponseCode.BAD_REQUEST);
                response.writeJsonFromObject(new JsonError(ResponseCode.BAD_REQUEST.getCode(), "SessionId Must Not Be Empty."));
                return;
            }
            serverSession = mcpSessionManager.getClientSeverSession(sessionId);
            if (serverSession == null) {
                response.setCode(ResponseCode.NOT_FOUND);
                response.writeJsonFromObject(new JsonError(ResponseCode.NOT_FOUND.getCode(), "Session Not Found."));
                return;
            }
        }
        String requestBody = request.getBodyAsString();
        if (logger.isLoggable(Level.FINE)) {
            logger.log(Level.FINE, "MCP Request Body: " + requestBody);
        }
        JsonMessage jsonMessage = getJsonMessage(requestBody);
        if (jsonMessage instanceof JsonResponse) {
            if (serverSession != null) {
                serverSession.getSseEmitter().completed();
                mcpSessionManager.removeClient(sessionId);
            }
            response.writeEmpty();
            return;
        }

        Object id = jsonMessage instanceof JsonRequest ? ((JsonRequest) jsonMessage).getId() : null;
        String method = jsonMessage instanceof JsonRequest ? ((JsonRequest) jsonMessage).getMethod() : ((JsonNotification) jsonMessage).getMethod();
        MethodHandler methodHandler = methodHandlerMap.get(method);
        if (methodHandler == null) {
            throw new McpException(new JsonError(ResponseCode.METHOD_NOT_SUPPORT.getCode(), "UnSupport MCP Method: " + method));
        }
        JsonResponse jsonResponse;
        try {
            jsonResponse = methodHandler.handle(serverSession, jsonMessage);
        } catch (McpException e) {
            jsonResponse = new JsonResponse();
            jsonResponse.setJsonrpc(Constants.JSONRPC);
            jsonResponse.setId(id);
            jsonResponse.setError(e.getJsonError());
        }
        if (jsonResponse == null) {
            response.writeEmpty();
            return;
        }
        McpRequestMetadata mcpRequestMetadata = new McpRequestMetadata();
        mcpRequestMetadata.setRequest(request);
        mcpRequestMetadata.setResponse(response);
        mcpRequestMetadata.setExecution(execution);
        mcpRequestMetadata.setMcpServerSession(serverSession);
        boolean result = doResponse(mcpRequestMetadata, jsonResponse);
        if (logger.isLoggable(Level.FINE)) {
            logger.log(Level.FINE, "MCP Response Body: " + JsonNullRemover.removeNullFields(getResponsiveConvert().convertTo(jsonResponse)));
        }
        if (!result) {
            response.writeEmpty();
        }
    }

    /**
     * 执行响应
     * @param mcpRequestMetadata MCP请求元数据
     * @param jsonResponse JSON响应
     * @return 是否已经对http请求响应
     */
    protected abstract boolean doResponse(McpRequestMetadata mcpRequestMetadata, JsonResponse jsonResponse);

    /**
     * 是否支持会话
     * @return 是否支持会话
     */
    protected boolean supportServerSession() {
        return true;
    }

    /**
     * 获取响应式转换器
     * @return 响应式转换器
     */
    protected ResponsiveConvert getResponsiveConvert() {
        return responsiveConvert;
    }

    /**
     * 获取JSON消息
     * @param requestBody 请求体内容
     * @return JSON消息
     */
    private JsonMessage getJsonMessage(String requestBody) {
        Map<String, Object> map = responsiveConvert.convertFromStringByType(requestBody, new TypeReference<Map<String, Object>>() {
        });
        if (map.containsKey("method") && map.containsKey("id")) {
            return converFromObject(responsiveConvert, map, JsonRequest.class);
        }
        else if (map.containsKey("method") && !map.containsKey("id")) {
            return converFromObject(responsiveConvert, map, JsonNotification.class);
        }
        else if (map.containsKey("result") || map.containsKey("error")) {
            return converFromObject(responsiveConvert, map, JsonResponse.class);
        }
        throw new IllegalArgumentException("UnSupport MCP JSON Message: " + requestBody);
    }

    /**
     * 初始化方法处理器
     */
    private void initMethodHandler() {
        methodHandlerMap.put("initialize", new InitializeMethodHandler(applicationContext, mcpServerCapability, environment));
        logger.info("MCP[Method], Method[initialize] Injected.");
        methodHandlerMap.put("notifications/initialized", new InitializedNotificationMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[notifications/initialized] Injected.");
        methodHandlerMap.put("ping", new PingMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[ping] Injected.");
        methodHandlerMap.put("tools/list", new ToolsListMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[tools/list] Injected.");
        methodHandlerMap.put("tools/call", new ToolsCallMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[tools/call] Injected.");
        methodHandlerMap.put("resources/list", new ResourcesListMethodHandler(applicationContext, mcpServerCapability, ResourceType.RESOURCE));
        logger.info("MCP[Method], Method[resources/list] Injected.");
        methodHandlerMap.put("resources/templates/list", new ResourcesListMethodHandler(applicationContext, mcpServerCapability, ResourceType.RESOURCE_TEMPLATE));
        logger.info("MCP[Method], Method[resources/templates/list] Injected.");
        methodHandlerMap.put("resources/read", new ResourcesReadMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[resources/read] Injected.");
        methodHandlerMap.put("prompts/list", new PromptsListMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[prompts/list] Injected.");
        methodHandlerMap.put("prompts/get", new PromptsGetMethodHandler(applicationContext, mcpServerCapability));
        logger.info("MCP[Method], Method[prompts/get] Injected.");
    }
}
